{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "73208366",
   "metadata": {},
   "source": [
    "# VM `execbuilder` 和估值"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "24628e27",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pytest\n",
    "\n",
    "import tvm\n",
    "from tvm import TVMError, relax\n",
    "from tvm.relax.testing.vm import check_saved_func\n",
    "from tvm.script import relax as R\n",
    "import tvm.testing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "836812a3",
   "metadata": {},
   "source": [
    "## `execute`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d715904c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ib = relax.ExecBuilder()\n",
    "with ib.function(\"func0\", num_inputs=2):\n",
    "    ib.emit_call(\"test.vm.add\", args=[ib.r(0), ib.r(1)], dst=ib.r(2))\n",
    "    ib.emit_ret(ib.r(2))\n",
    "ex = ib.get()\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "a = tvm.nd.array(\n",
    "    np.random.rand(\n",
    "        4,\n",
    "    )\n",
    ")\n",
    "b = tvm.nd.array(\n",
    "    np.random.rand(\n",
    "        4,\n",
    "    )\n",
    ")\n",
    "\n",
    "add_res = check_saved_func(vm, \"func0\", a, b)\n",
    "tvm.testing.assert_allclose(add_res.numpy(), a.numpy() + b.numpy(), rtol=1e-7, atol=1e-7)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aba254b7",
   "metadata": {},
   "source": [
    "## 多函数测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "84bea8d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "ib = relax.ExecBuilder()\n",
    "with ib.function(\"func0\", num_inputs=2):\n",
    "    ib.emit_call(\"test.vm.add\", args=[ib.r(0), ib.r(1)], dst=ib.r(2))\n",
    "    ib.emit_ret(ib.r(2))\n",
    "with ib.function(\"func1\", num_inputs=2):\n",
    "    ib.emit_call(\"test.vm.mul\", args=[ib.r(0), ib.r(1)], dst=ib.r(2))\n",
    "    ib.emit_ret(ib.r(2))\n",
    "ex = ib.get()\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "a = tvm.nd.array(\n",
    "    np.random.rand(\n",
    "        4,\n",
    "    )\n",
    ")\n",
    "b = tvm.nd.array(\n",
    "    np.random.rand(\n",
    "        4,\n",
    "    )\n",
    ")\n",
    "mul_res = check_saved_func(vm, \"func1\", a, b)\n",
    "add_res = check_saved_func(vm, \"func0\", a, b)\n",
    "tvm.testing.assert_allclose(add_res.numpy(), a.numpy() + b.numpy(), rtol=1e-7, atol=1e-7)\n",
    "tvm.testing.assert_allclose(mul_res.numpy(), a.numpy() * b.numpy(), rtol=1e-7, atol=1e-7)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "701c9c0a",
   "metadata": {},
   "source": [
    "## VM 检查"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b6c0113e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ib = relax.ExecBuilder()\n",
    "with pytest.raises(TVMError):\n",
    "    with ib.function(\"func0\", num_inputs=2):\n",
    "        ib.emit_call(\"test.vm.add\", args=[ib.r(0), ib.r(2)], dst=ib.r(2))\n",
    "        ib.emit_ret(ib.r(2))\n",
    "    ib.get()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14eec0b4",
   "metadata": {},
   "source": [
    "## 负常数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7212191d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ib = relax.ExecBuilder()\n",
    "\n",
    "with ib.function(\"func0\", num_inputs=1):\n",
    "    ib.emit_call(\"test.vm.add_scalar\", args=[ib.imm(-3), ib.r(0)], dst=ib.r(1))\n",
    "    ib.emit_ret(ib.r(1))\n",
    "ib.get()\n",
    "\n",
    "ex = ib.get()\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "assert vm[\"func0\"](1) == -2\n",
    "assert vm[\"func0\"](-3) == -6"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e345775e",
   "metadata": {},
   "source": [
    "## 缓存"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2b98fad3",
   "metadata": {},
   "outputs": [],
   "source": [
    "ib = relax.ExecBuilder()\n",
    "\n",
    "with ib.function(\"func0\", num_inputs=1):\n",
    "    x0 = ib.convert_constant(\"str0\")\n",
    "    x1 = ib.convert_constant(\"str0\")\n",
    "    # cache constant str\n",
    "    assert x0 == x1\n",
    "    s0 = ib.convert_constant(tvm.runtime.container.ShapeTuple([1, 2]))\n",
    "    s1 = ib.convert_constant(tvm.runtime.container.ShapeTuple([1, 2]))\n",
    "    s2 = ib.convert_constant(tvm.runtime.container.ShapeTuple([1, 3]))\n",
    "    assert s0 == s1\n",
    "    assert s1 != s2\n",
    "    y0 = ib.convert_constant(tvm.nd.array(np.array([1, 2, 3]).astype(\"int32\")))\n",
    "    y1 = ib.convert_constant(tvm.nd.array(np.array([1, 2, 3]).astype(\"int32\")))\n",
    "    assert y0 == y1\n",
    "    ib.emit_ret(ib.r(0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "224d3aaf",
   "metadata": {},
   "source": [
    "## VM 格式化\n",
    "\n",
    "formalize:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b2122520",
   "metadata": {},
   "outputs": [],
   "source": [
    "ib0 = relax.ExecBuilder()\n",
    "ib1 = relax.ExecBuilder()\n",
    "with ib0.function(\"func0\", num_inputs=2):\n",
    "    ib0.emit_call(\"test.vm.add\", args=[ib0.r(0), ib0.r(1)], dst=ib0.r(100))\n",
    "    ib0.emit_call(\"test.vm.mul\", args=[ib0.r(1), ib0.r(100)], dst=ib0.r(50))\n",
    "    ib0.emit_ret(ib0.r(50))\n",
    "with ib1.function(\"func0\", num_inputs=2):\n",
    "    ib1.emit_call(\"test.vm.add\", args=[ib1.r(0), ib1.r(1)], dst=ib1.r(2))\n",
    "    ib1.emit_call(\"test.vm.mul\", args=[ib1.r(1), ib1.r(2)], dst=ib1.r(3))\n",
    "    ib1.emit_ret(ib1.r(3))\n",
    "exec0 = ib0.get()\n",
    "exec1 = ib1.get()\n",
    "assert exec0.as_text() == exec1.as_text()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9066c5bb",
   "metadata": {},
   "source": [
    "## operand"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "960a4c39",
   "metadata": {},
   "outputs": [],
   "source": [
    "ib0 = relax.ExecBuilder()\n",
    "with ib0.function(\"func0\", num_inputs=2):\n",
    "    ib0.emit_call(\"test.vm.add_scalar\", args=[ib0.r(0), ib0.r(1)], dst=ib0.r(2))\n",
    "    ib0.emit_ret(ib0.r(2))\n",
    "exec0 = ib0.get()\n",
    "vm = relax.VirtualMachine(exec0, tvm.cpu())\n",
    "res = vm[\"func0\"](2, 3)\n",
    "assert res == 5\n",
    "\n",
    "ib1 = relax.ExecBuilder()\n",
    "with ib1.function(\"func1\", num_inputs=1):\n",
    "    ib1.emit_call(\"test.vm.get_device_id\", args=[ib1.r(0)], dst=ib1.r(1))\n",
    "    ib1.emit_ret(ib1.r(1))\n",
    "exec1 = ib1.get()\n",
    "vm = relax.VirtualMachine(exec1, tvm.cpu())\n",
    "res = vm[\"func1\"](tvm.cpu(3))\n",
    "assert res == 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "791dc5e0",
   "metadata": {},
   "source": [
    "## shapeof"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c36bfa17",
   "metadata": {},
   "outputs": [],
   "source": [
    "ib = relax.ExecBuilder()\n",
    "shape = (32, 16)\n",
    "arr = tvm.nd.array(np.random.rand(*shape))\n",
    "with ib.function(\"main\", num_inputs=0):\n",
    "    ib.emit_call(\"vm.builtin.shape_of\", args=[arr], dst=ib.r(0))\n",
    "    ib.emit_ret(ib.r(0))\n",
    "ex = ib.get()\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "res = vm[\"main\"]()\n",
    "for i, s in enumerate(res):\n",
    "    assert s == shape[i]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bca97b9",
   "metadata": {},
   "source": [
    "## storage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "877cf3fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "dtype = tvm.DataType(\"float32\")\n",
    "shape = (4, 6)\n",
    "ib = relax.ExecBuilder()\n",
    "with ib.function(\"main\", num_inputs=0):\n",
    "    ib.emit_call(\n",
    "        \"vm.builtin.alloc_storage\",\n",
    "        args=[\n",
    "            ib.vm_state(),\n",
    "            (24,),\n",
    "            ib.convert_constant(0),\n",
    "            dtype,\n",
    "            ib.convert_constant(\"global\"),\n",
    "        ],\n",
    "        dst=ib.r(1),\n",
    "    )\n",
    "    ib.emit_call(\n",
    "        \"vm.builtin.alloc_tensor\", args=[ib.r(1), ib.imm(0), shape, dtype], dst=ib.r(2)\n",
    "    )\n",
    "    ib.emit_ret(ib.r(2))\n",
    "ex = ib.get()\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "res = vm[\"main\"]()\n",
    "assert res.device == tvm.cpu()\n",
    "assert res.shape == shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e0282ff",
   "metadata": {},
   "source": [
    "## goto"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "121e4c98",
   "metadata": {},
   "outputs": [],
   "source": [
    "ib = relax.ExecBuilder()\n",
    "with ib.function(\"main\", num_inputs=2):\n",
    "    ib.emit_call(\"test.vm.add\", args=[ib.r(0), ib.r(1)], dst=ib.r(2))\n",
    "    ib.emit_goto(2)\n",
    "    ib.emit_call(\"test.vm.mul\", args=[ib.r(2), ib.r(1)], dst=ib.r(2))\n",
    "    ib.emit_ret(ib.r(2))\n",
    "ex = ib.get()\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "a = tvm.nd.array(\n",
    "    np.random.rand(\n",
    "        4,\n",
    "    )\n",
    ")\n",
    "b = tvm.nd.array(\n",
    "    np.random.rand(\n",
    "        4,\n",
    "    )\n",
    ")\n",
    "res = check_saved_func(vm, \"main\", a, b)\n",
    "tvm.testing.assert_allclose(res.numpy(), a.numpy() + b.numpy(), rtol=1e-7, atol=1e-7)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db5ce72f",
   "metadata": {},
   "source": [
    "## `if`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "af1351b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "ib = relax.ExecBuilder()\n",
    "with ib.function(\"main\", num_inputs=3):\n",
    "    ib.emit_if(ib.r(0), 3)\n",
    "    ib.emit_call(\"test.vm.add\", args=[ib.r(1), ib.r(2)], dst=ib.r(3))\n",
    "    ib.emit_goto(2)\n",
    "    ib.emit_call(\"test.vm.mul\", args=[ib.r(1), ib.r(2)], dst=ib.r(3))\n",
    "    ib.emit_ret(ib.r(3))\n",
    "ex = ib.get()\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "a = tvm.nd.array(\n",
    "    np.random.rand(\n",
    "        4,\n",
    "    )\n",
    ")\n",
    "b = tvm.nd.array(\n",
    "    np.random.rand(\n",
    "        4,\n",
    "    )\n",
    ")\n",
    "res = vm[\"main\"](0, a, b)\n",
    "tvm.testing.assert_allclose(res.numpy(), a.numpy() * b.numpy(), rtol=1e-7, atol=1e-7)\n",
    "res = vm[\"main\"](1, a, b)\n",
    "tvm.testing.assert_allclose(res.numpy(), a.numpy() + b.numpy(), rtol=1e-7, atol=1e-7)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d9dac09",
   "metadata": {},
   "source": [
    "## 闭包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9bb58616",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_1447999/4112004499.py:19: UserWarning: Returning type `relax.vm.Closure` which is not registered via register_object, fallback to Object\n",
      "  clo = vm[\"main\"](w_inp, x_inp)\n"
     ]
    }
   ],
   "source": [
    "ib = relax.ExecBuilder()\n",
    "with ib.function(\"lifted_func_1\", num_inputs=4):\n",
    "    ib.emit_call(\"test.vm.add\", args=[ib.r(0), ib.r(1)], dst=ib.r(4))\n",
    "    ib.emit_call(\"test.vm.add\", args=[ib.r(2), ib.r(4)], dst=ib.r(5))\n",
    "    ib.emit_call(\"test.vm.add\", args=[ib.r(3), ib.r(5)], dst=ib.r(6))\n",
    "    ib.emit_ret(ib.r(6))\n",
    "with ib.function(\"main\", num_inputs=2):\n",
    "    ib.emit_call(\n",
    "        \"vm.builtin.make_closure\", args=[ib.f(\"lifted_func_1\"), ib.r(0), ib.r(1)], dst=ib.r(2)\n",
    "    )\n",
    "    ib.emit_ret(ib.r(2))\n",
    "\n",
    "ex = ib.get()\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "w_inp = tvm.nd.array(np.random.rand(2, 3))\n",
    "x_inp = tvm.nd.array(np.random.rand(2, 3))\n",
    "y_inp = tvm.nd.array([[3.1, 4.0, 5.0], [6.0, 7.1, 9.0]])\n",
    "z_inp = tvm.nd.array(np.random.rand(2, 3))\n",
    "clo = vm[\"main\"](w_inp, x_inp)\n",
    "res = vm.invoke_closure(clo, y_inp, z_inp)\n",
    "tvm.testing.assert_allclose(\n",
    "    res.numpy(), w_inp.numpy() + x_inp.numpy() + y_inp.numpy() + z_inp.numpy()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "873cd5df",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
